import sys
import os.path

from PyQt4 import QtCore, QtGui
QtCore.Signal = QtCore.pyqtSignal

import vtk
from vtk.qt4.QVTKRenderWindowInteractor import QVTKRenderWindowInteractor

import numpy as np

from SPNAnnotation import SPNAnnotation

class VTKFrame(QtGui.QFrame):
    def __init__(self, parent = None):
        super(VTKFrame, self).__init__(parent)

        self.vtkWidget = QVTKRenderWindowInteractor(self)
        vl = QtGui.QVBoxLayout(self)
        vl.addWidget(self.vtkWidget)
        vl.setContentsMargins(0, 0, 0, 0)
 
        self.ren = vtk.vtkRenderer()
        self.backgroundColor = (0,0,0)
        self.ren.SetBackground(self.backgroundColor)
        self.vtkWidget.GetRenderWindow().AddRenderer(self.ren)
        self.iren = self.vtkWidget.GetRenderWindow().GetInteractor()
        self.iren.SetInteractorStyle(vtk.vtkInteractorStyleTrackballCamera())
 
        self.metaReader=vtk.vtkMetaImageReader() 
        self.accum=vtk.vtkImageAccumulate()        
        self.accum.SetInputConnection(self.metaReader.GetOutputPort())

        self.annoExtractVOI=vtk.vtkExtractVOI()
        self.annoExtractVOI.SetInputConnection(self.metaReader.GetOutputPort())

        self.threshold = vtk.vtkImageThreshold()
        self.threshold.ReplaceInOff()
        self.threshold.ReplaceOutOff()
        self.threshold.SetInputConnection(self.metaReader.GetOutputPort())

        # Create transfer mapping scalar value to opacity
        self.opacityTransferFunction = vtk.vtkPiecewiseFunction()
        # Create transfer mapping scalar value to color
        self.colorTransferFunction = vtk.vtkColorTransferFunction()

        self.volumeProperty = vtk.vtkVolumeProperty()
        self.volumeProperty.SetColor(self.colorTransferFunction)
        self.volumeProperty.SetScalarOpacity(self.opacityTransferFunction)
        self.volumeProperty.ShadeOn()
        self.volumeProperty.SetInterpolationTypeToLinear()
        
        self.volumeMapper = vtk.vtkFixedPointVolumeRayCastMapper()
        # volumeMapper.SetVolumeRayCastFunction(compositeFunction)
        self.volumeMapper.SetInputConnection(self.threshold.GetOutputPort())

        # The volume holds the mapper and the property and
        # can be used to position/orient the volume
        self.volume = vtk.vtkVolume()
        self.volume.SetMapper(self.volumeMapper)
        self.volume.SetProperty(self.volumeProperty)

        self.ren.AddVolume(self.volume)
        self.volume.Update()
        self.ren.ResetCamera()

        self.minThreshold = 0
        self.maxThreshold = 0
        self.curMinVal = 0
        self.curMaxVal = 0
        self.curAnnotation = None
        self.cropSize = 1

        self.sphereWidget = vtk.vtkSphereWidget()
        self.sphereWidget.SetInteractor(self.iren)
        self.sphereWidget.Off()
        self._initialized = False

    def showEvent(self, evt):
        if not self._initialized:
            self.iren.Initialize()
            self.startTimer(30)
            self._initialized = True

    def timerEvent(self, evt):
        self.vtkWidget.GetRenderWindow().Render()

    def setCTData(self, fileName):        
        self.metaReader.SetFileName(fileName)
        self.accum.Update()    
            
        self.minThreshold = self.accum.GetMin()[0]
        self.maxThreshold = self.accum.GetMax()[0]
        self.ren.ResetCamera()

    def setThreshold(self, minVal, maxVal):
        self.curMinVal = minVal
        self.curMaxVal = maxVal
        self.threshold.ThresholdBetween(minVal, maxVal)
        self.threshold.Update()
        
        self.opacityTransferFunction.RemoveAllPoints()
        self.opacityTransferFunction.AddPoint(minVal, 0.0)
        self.opacityTransferFunction.AddPoint(maxVal, 1.0)

        self.colorTransferFunction.RemoveAllPoints()
        # self.colorTransferFunction.AddRGBPoint(minVal, 0.0, 0.0, 0.0)
        # self.colorTransferFunction.AddRGBPoint(minVal + (maxVal - minVal) / 4, 1.0, 0.0, 0.0)
        # self.colorTransferFunction.AddRGBPoint(minVal + 2 * (maxVal - minVal) / 4, 0.0, 0.0, 1.0)
        # self.colorTransferFunction.AddRGBPoint(minVal + 3 * (maxVal - minVal) / 4, 0.0, 1.0, 0.0)
        # self.colorTransferFunction.AddRGBPoint(maxVal, 0.0, 0.2, 0.0)
        self.colorTransferFunction.AddRGBPoint(minVal, self.backgroundColor[0], self.backgroundColor[1], self.backgroundColor[2])
        self.colorTransferFunction.AddRGBPoint(maxVal, 1.0 - self.backgroundColor[0], 1.0 - self.backgroundColor[1], 1.0 - self.backgroundColor[2])
        
    def setAnnotation(self, anno):
        imageData = self.metaReader.GetOutput()
        origin = np.array(imageData.GetOrigin())
        self.curAnnotation = anno     
        coord = np.float32(np.array(anno.coord))
        # coord = coord - origin
        diameter = np.float(anno.diameter)
        self.sphereWidget.PlaceWidget(coord[0] - diameter, coord[0] + diameter, coord[1] - diameter, coord[1] + diameter, coord[2] - diameter, coord[2] + diameter)      
        self.sphereWidget.On()
        self.sphereWidget.ScaleOff()
        self.sphereWidget.TranslationOff()

    def toggleAnnotationMode(self, enabled):
        if enabled and self.curAnnotation != None:
            coord = np.float32(np.array(self.curAnnotation.coord))
            diameter = np.float(self.curAnnotation.diameter)

            imageData = self.metaReader.GetOutput()
            space = np.array(imageData.GetSpacing())
            origin = np.array(imageData.GetOrigin())
            center = (coord - origin) / space
            size = np.array((diameter, diameter, diameter)) / space * self.cropSize

            self.annoExtractVOI.SetVOI(int(center[0] - size[0]), int(center[0] + size[0]),
                                        int(center[1] - size[1]), int(center[1] + size[1]),
                                        int(center[2] - size[2]), int(center[2] + size[2]))

            self.accum.SetInputConnection(self.annoExtractVOI.GetOutputPort())
            self.accum.Update()
            self.curMinVal = self.accum.GetMin()[0]
            self.curMaxVal = self.accum.GetMax()[0]
            self.threshold.SetInputConnection(self.annoExtractVOI.GetOutputPort())
            self.threshold.ThresholdBetween(self.curMinVal, self.curMaxVal)
            return self.curMinVal, self.curMaxVal
        else:
            self.threshold.SetInputConnection(self.metaReader.GetOutputPort())
            self.threshold.ThresholdBetween(self.minThreshold, self.maxThreshold)
            self.showMarker(False)
            return self.minThreshold, self.maxThreshold

    def showMarker(self, show):
        if show:
            self.sphereWidget.On()
        else:
            self.sphereWidget.Off()

    def invertBackgroundColor(self):
        self.backgroundColor = (1.0 - self.backgroundColor[0], 1.0 - self.backgroundColor[1], 1.0 - self.backgroundColor[2])
        self.ren.SetBackground(self.backgroundColor)
        self.colorTransferFunction.RemoveAllPoints()
        self.colorTransferFunction.AddRGBPoint(self.curMinVal, self.backgroundColor[0], self.backgroundColor[1], self.backgroundColor[2])
        self.colorTransferFunction.AddRGBPoint(self.curMaxVal, 1.0 - self.backgroundColor[0], 1.0 - self.backgroundColor[1], 1.0 - self.backgroundColor[2])
        self.update()

    def update(self):
        self.volume.Update()
